import pandas as pd
import mysql.connector
from mysql.connector import Error


class MySQLDatabase:
    def __init__(self, host, user, password, database):
        self.config = {
            'user': user,
            'password': password,
            'host': host,
            'database': database
        }
        self.connection = None
        self.cursor = None

    def open(self):
        """建立数据库连接"""
        try:
            self.connection = mysql.connector.connect(**self.config)
            if self.connection.is_connected():
                print("成功连接到数据库")
                self.cursor = self.connection.cursor()
        except Error as e:
            print(f"连接数据库失败: {e}")

    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        print("数据库连接已关闭")

    def insert_data(self, df, table_name, selected_columns):
        """将DataFrame中的数据插入到MySQL表中
        df: DataFrame对象
        table_name: 表名
        selected_columns: 用户指定的列名列表
        """
        # 验证所选列是否存在于DataFrame中
        selected_columns = [col for col in selected_columns if col in df.columns.tolist()]

        # 动态生成INSERT语句
        columns_placeholder = ', '.join(selected_columns)
        placeholders = ', '.join(['%s'] * len(selected_columns))
        sql = f"INSERT INTO {table_name} ({columns_placeholder}) VALUES ({placeholders})"

        for index, row in df.iterrows():
            # 只提取用户指定的列数据
            data_to_insert = [row[col] if not pd.isna(row[col]) else None for col in selected_columns]
            self.cursor.execute(sql, data_to_insert)

        # 提交事务
        self.connection.commit()
        print(f"成功插入{len(df)}条数据")

    def execute_query(self, query, params=None):
        """执行单个查询"""
        try:
            self.cursor.execute(query, params)
            self.connection.commit()
            print("查询成功")
        except Error as e:
            print(f"执行查询失败: {e}")

    def fetch_all(self, query, params=None):
        """执行查询并返回所有结果"""
        try:
            self.cursor.execute(query, params)
            results = self.cursor.fetchall()
            return results
        except Error as e:
            print(f"获取数据失败: {e}")
